import os
import re
import shutil

# 需要处理的md文件
md_files = [
    r'd:\笔记\安全渗透\SQL注入.md',
    r'd:\笔记\安全渗透\SQL注入绕过.md',
    r'd:\笔记\安全渗透\Javascript.md'
]

# 匹配图片路径的正则
img_pattern = re.compile(r'!\[.*?\]\((.*?)\)')

for md_file in md_files:
    md_dir = os.path.dirname(md_file)
    md_name = os.path.splitext(os.path.basename(md_file))[0]
    img_dir = os.path.join(md_dir, f"{md_name}")
    os.makedirs(img_dir, exist_ok=True)
    img_map = {}  # 原路径: 新路径
    img_count = 1

    # 读取md内容
    with open(md_file, 'r', encoding='utf-8') as f:
        content = f.read()
        matches = img_pattern.findall(content)
        for img_path in matches:
            # 只处理本地图片
            if img_path.startswith('http'):
                continue
            # 处理相对路径，兼容 \ 和 /
            norm_img_path = img_path.replace('/', os.sep).replace('\\', os.sep)
            # 如果是以 data\static\ 开头，拼接到笔记目录下
            if norm_img_path.startswith('data' + os.sep):
                abs_img_path = os.path.join(md_dir, norm_img_path)
            else:
                abs_img_path = os.path.join(md_dir, norm_img_path)
            if not os.path.isfile(abs_img_path):
                continue
            ext = os.path.splitext(abs_img_path)[1]
            new_img_name = f'img_{img_count:03d}{ext}'
            new_img_path = os.path.join(img_dir, new_img_name)
            if img_path not in img_map:
                shutil.copy2(abs_img_path, new_img_path)
                # 新路径为相对md文件的路径
                rel_new_img_path = os.path.relpath(new_img_path, md_dir).replace('\\', '/')
                img_map[img_path] = rel_new_img_path
                img_count += 1

    # 替换md内容中的图片路径
    def replace_img(m):
        old = m.group(1)
        return m.group(0).replace(old, img_map.get(old, old))
    new_content = img_pattern.sub(replace_img, content)

    # 写回md文件
    with open(md_file, 'w', encoding='utf-8') as f:
        f.write(new_content)

    # 输出映射表
    print(f"{md_file} 图片路径映射：")
    for k, v in img_map.items():
        print(f"{k} -> {v}")
